import pygame
import sys
import time
from typing import List, Tuple, Dict, Set
import heapq
from config import *

class Boss:
    def __init__(self, name="恶龙"):
        self.name = name
        self.health = 50
        self.attack = 8
        self.defense = 2
        
    def attack_player(self, player):
        """BOSS攻击玩家"""
        damage = max(0, self.attack)
        player.health -= damage
        return damage

class Player:
    def __init__(self, initial_health=100, initial_mp=30):
        self.health = initial_health
        self.mp = initial_mp
        self.max_health = initial_health
        self.max_mp = initial_mp
        # 技能冷却: [普攻冷却, 大招冷却, 治疗冷却]
        self.cooldowns = [0, 0, 0]
        
    def regenerate_mp(self):
        """每回合恢复2点MP"""
        self.mp = min(self.max_mp, self.mp + 2)
        
    def update_cooldowns(self):
        """更新技能冷却时间"""
        for i in range(len(self.cooldowns)):
            if self.cooldowns[i] > 0:
                self.cooldowns[i] -= 1
                
    def use_skill(self, skill_id, boss):
        """
        使用技能攻击BOSS或治疗自己
        skill_id: 0=普攻, 1=大招, 2=治疗
        """
        if self.cooldowns[skill_id] > 0:
            return f"技能冷却中，还需{self.cooldowns[skill_id]}回合"
        
        if skill_id == 0:  # 普通攻击
            damage = 5
            mp_cost = 0
            cooldown = 0
            boss.health -= damage
            result = f"普通攻击造成{damage}点伤害!"
        elif skill_id == 1:  # 大招
            if self.mp < 8:
                return "MP不足，无法使用大招!"
            damage = 12
            mp_cost = 8
            cooldown = 2
            boss.health -= damage
            result = f"大招造成{damage}点伤害!"
        elif skill_id == 2:  # 治疗
            if self.mp < 6:
                return "MP不足，无法使用治疗!"
            damage = 0
            mp_cost = 6
            cooldown = 1
            heal_amount = 10
            self.health = min(self.max_health, self.health + heal_amount)
            result = f"治疗恢复{heal_amount}点生命值!"
        else:
            return "无效的技能ID"
            
        self.mp -= mp_cost
        self.cooldowns[skill_id] = cooldown
        return result

class BossVisualization:
    def __init__(self, boss, player):
        # 初始化pygame
        pygame.init()
        self.screen_width = 800
        self.screen_height = 600
        self.screen = pygame.display.set_mode((self.screen_width, self.screen_height))
        pygame.display.set_caption('BOSS战斗')
        
        self.boss = boss
        self.player = player
        self.font_large = pygame.font.SysFont('kaiti', 36)
        self.font_medium = pygame.font.SysFont('kaiti', 24)
        self.font_small = pygame.font.SysFont('kaiti', 18)
        
        # 技能信息
        self.skills = [
            {"name": "普通攻击", "damage": "5", "mp_cost": "0", "cooldown": "0"},
            {"name": "大招", "damage": "12", "mp_cost": "8", "cooldown": "2"},
            {"name": "治疗术", "heal": "10", "mp_cost": "6", "cooldown": "1"}
        ]
        
        # 战斗日志
        self.battle_log = ["战斗开始!"]
        self.max_logs = 5
        
        # 推荐技能序列
        self.recommended_sequence = []
        self.current_recommendation = None

    def display_boss(self):
        """显示BOSS信息"""
        self.screen.fill((30, 30, 50))  # 深蓝色背景
        
        # 绘制BOSS区域
        boss_area = pygame.Rect(50, 50, 700, 200)
        pygame.draw.rect(self.screen, (60, 30, 30), boss_area, border_radius=10)
        
        # BOSS名称
        boss_name = self.font_large.render(f"{self.boss.name}", True, COLOR_WHITE)
        self.screen.blit(boss_name, (400 - boss_name.get_width()//2, 70))
        
        # BOSS血量条
        boss_hp_bg = pygame.Rect(150, 120, 500, 30)
        boss_hp = pygame.Rect(150, 120, 500 * (self.boss.health / 50), 30)
        pygame.draw.rect(self.screen, (100, 0, 0), boss_hp_bg)  # 血条背景
        pygame.draw.rect(self.screen, (255, 0, 0), boss_hp)     # 实际血量
        boss_hp_text = self.font_medium.render(f"HP: {self.boss.health}/50", True, COLOR_WHITE)
        self.screen.blit(boss_hp_text, (400 - boss_hp_text.get_width()//2, 125))
        
        # 攻击力
        boss_atk = self.font_small.render(f"攻击: {self.boss.attack}", True, COLOR_WHITE)
        self.screen.blit(boss_atk, (150, 170))
        
        # 防御力
        boss_def = self.font_small.render(f"防御: {self.boss.defense}", True, COLOR_WHITE)
        self.screen.blit(boss_def, (250, 170))
    
    def display_player(self):
        """显示玩家信息"""
        # 绘制玩家区域
        player_area = pygame.Rect(50, 300, 700, 150)
        pygame.draw.rect(self.screen, (30, 60, 30), player_area, border_radius=10)
        
        # 玩家名称
        player_name = self.font_medium.render("勇者", True, COLOR_WHITE)
        self.screen.blit(player_name, (400 - player_name.get_width()//2, 310))
        
        # 血量条
        hp_bg = pygame.Rect(150, 350, 500, 20)
        hp = pygame.Rect(150, 350, 500 * (self.player.health / self.player.max_health), 20)
        pygame.draw.rect(self.screen, (100, 0, 0), hp_bg)  # 血条背景
        pygame.draw.rect(self.screen, (255, 0, 0), hp)     # 实际血量
        hp_text = self.font_small.render(f"HP: {self.player.health}/{self.player.max_health}", True, COLOR_WHITE)
        self.screen.blit(hp_text, (400 - hp_text.get_width()//2, 350))
        
        # MP条
        mp_bg = pygame.Rect(150, 380, 500, 20)
        mp = pygame.Rect(150, 380, 500 * (self.player.mp / self.player.max_mp), 20)
        pygame.draw.rect(self.screen, (30, 30, 150), mp_bg)  # MP条背景
        pygame.draw.rect(self.screen, (0, 0, 255), mp)      # 实际MP
        mp_text = self.font_small.render(f"MP: {self.player.mp}/{self.player.max_mp}", True, COLOR_WHITE)
        self.screen.blit(mp_text, (400 - mp_text.get_width()//2, 380))
        
        # 技能冷却
        cooldown_text = self.font_small.render(
            f"技能冷却: 普攻({self.player.cooldowns[0]}), 大招({self.player.cooldowns[1]}), 治疗({self.player.cooldowns[2]})", 
            True, COLOR_WHITE
        )
        self.screen.blit(cooldown_text, (150, 410))
    
    def display_skills(self):
        """显示技能按钮"""
        skill_area = pygame.Rect(50, 470, 700, 80)
        pygame.draw.rect(self.screen, (50, 50, 60), skill_area, border_radius=10)
        
        btn_width = 200
        btn_height = 60
        btn_spacing = 25
        
        for i, skill in enumerate(self.skills):
            x = 100 + i * (btn_width + btn_spacing)
            y = 480
            
            # 确定按钮颜色(可用为绿色，冷却中为灰色，MP不足为红色)
            if self.player.cooldowns[i] > 0:
                color = (100, 100, 100)  # 灰色
            elif (i == 1 and self.player.mp < 8) or (i == 2 and self.player.mp < 6):
                color = (150, 50, 50)    # 红色
            else:
                color = (50, 150, 50)    # 绿色
                
            # 绘制按钮
            btn_rect = pygame.Rect(x, y, btn_width, btn_height)
            pygame.draw.rect(self.screen, color, btn_rect, border_radius=5)
            
            # 技能名称
            skill_name = self.font_medium.render(skill["name"], True, COLOR_WHITE)
            self.screen.blit(skill_name, (x + btn_width//2 - skill_name.get_width()//2, y + 5))
            
            # 技能信息
            if "damage" in skill:
                damage_text = self.font_small.render(f"伤害: {skill['damage']}", True, COLOR_WHITE)
                self.screen.blit(damage_text, (x + 10, y + 35))
            elif "heal" in skill:
                heal_text = self.font_small.render(f"恢复: {skill['heal']}", True, COLOR_WHITE)
                self.screen.blit(heal_text, (x + 10, y + 35))
                
            mp_text = self.font_small.render(f"MP: {skill['mp_cost']}", True, COLOR_WHITE)
            self.screen.blit(mp_text, (x + btn_width - 60, y + 35))
    
    def display_battle_log(self):
        """显示战斗日志"""
        log_area = pygame.Rect(50, 560, 700, 30)
        pygame.draw.rect(self.screen, (30, 30, 30), log_area)
        
        if self.battle_log:
            latest_log = self.battle_log[-1]
            log_text = self.font_small.render(latest_log, True, COLOR_WHITE)
            self.screen.blit(log_text, (60, 565))
    
    def add_battle_log(self, message):
        """添加战斗日志"""
        self.battle_log.append(message)
        if len(self.battle_log) > self.max_logs:
            self.battle_log.pop(0)
    
    def display_recommendation(self):
        """显示推荐技能"""
        if self.current_recommendation is not None:
            rec_text = self.font_medium.render(
                f"推荐技能: {self.skills[self.current_recommendation]['name']}", 
                True, (255, 215, 0)
            )
            self.screen.blit(rec_text, (50, 260))
    
    def update_display(self):
        """更新显示"""
        self.display_boss()
        self.display_player()
        self.display_skills()
        self.display_battle_log()
        self.display_recommendation()
        pygame.display.flip()
    
    def handle_events(self):
        """处理事件"""
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            
            if event.type == pygame.MOUSEBUTTONDOWN:
                mouse_pos = pygame.mouse.get_pos()
                
                # 检查技能按钮点击
                btn_width = 200
                btn_height = 60
                btn_spacing = 25
                
                for i in range(3):
                    x = 100 + i * (btn_width + btn_spacing)
                    y = 480
                    btn_rect = pygame.Rect(x, y, btn_width, btn_height)
                    
                    if btn_rect.collidepoint(mouse_pos):
                        return i  # 返回点击的技能ID
        
        return None  # 没有点击技能

class BranchAndBoundSolver:
    """分支限界法求解器"""
    def __init__(self, boss_hp=50, player_hp=100, player_mp=30):
        self.initial_boss_hp = boss_hp
        self.initial_player_hp = player_hp
        self.initial_player_mp = player_mp
        self.max_turns = 20
        self.best_sequence = None
        self.best_score = float('inf')  # 回合数越少越好
    
    def estimate_remaining_turns(self, boss_hp, player_mp, cooldowns):
        """估计至少还需要多少回合才能击败BOSS"""
        if boss_hp <= 0:
            return 0
            
        # 考虑MP和冷却的情况下，预估平均每回合可造成的伤害
        avg_damage_per_turn = 5  # 至少可以使用普攻
        
        # 如果有足够的MP和没有冷却，可以考虑使用大招
        if player_mp >= 8 and cooldowns[1] == 0:
            avg_damage_per_turn = max(avg_damage_per_turn, 12/3)  # 每3回合使用一次大招
        
        return max(1, (boss_hp + avg_damage_per_turn - 1) // avg_damage_per_turn)
    
    def calculate_score(self, turns_used, player_hp, player_mp):
        """计算得分"""
        return 1000 - 50 * turns_used + 2 * player_hp + player_mp
    
    def solve(self):
        """使用分支限界法求解最优技能序列"""
        # 优先队列，按照预估总回合数排序
        priority_queue = []
        
        # 初始状态: (预估总回合数, 已用回合数, BOSS血量, 玩家血量, 玩家MP, 技能冷却, 已选择的技能序列)
        initial_cooldowns = [0, 0, 0]
        heapq.heappush(priority_queue, (
            self.estimate_remaining_turns(self.initial_boss_hp, self.initial_player_mp, initial_cooldowns),
            0,  # 已用回合数
            self.initial_boss_hp,
            self.initial_player_hp,
            self.initial_player_mp,
            initial_cooldowns,
            []  # 技能序列
        ))
        
        while priority_queue:
            estimated_turns, turns_used, boss_hp, player_hp, player_mp, cooldowns, skill_sequence = heapq.heappop(priority_queue)
            
            # 剪枝: 已用回合数超过最优解或超过最大回合限制
            if turns_used >= self.best_score or turns_used >= self.max_turns:
                continue
            
            # 胜利条件: BOSS血量为0且玩家存活
            if boss_hp <= 0 and player_hp > 0:
                score = self.calculate_score(turns_used, player_hp, player_mp)
                if turns_used < self.best_score:
                    self.best_score = turns_used
                    self.best_sequence = skill_sequence
                continue
            
            # 失败条件: 玩家血量为0
            if player_hp <= 0:
                continue
            
            # BOSS攻击玩家
            new_player_hp = player_hp - 8  # BOSS每回合造成8点伤害
            
            if new_player_hp <= 0:
                continue  # 玩家会死亡，跳过
            
            # 玩家MP回复
            new_player_mp = min(30, player_mp + 2)
            
            # 更新技能冷却
            new_cooldowns = cooldowns.copy()
            for i in range(len(new_cooldowns)):
                if new_cooldowns[i] > 0:
                    new_cooldowns[i] -= 1
            
            # 尝试三种技能
            for skill_id in range(3):
                # 检查冷却和MP
                if new_cooldowns[skill_id] > 0:
                    continue
                
                skill_mp_cost = [0, 8, 6][skill_id]
                if new_player_mp < skill_mp_cost:
                    continue
                
                # 使用技能
                skill_new_player_mp = new_player_mp - skill_mp_cost
                skill_boss_hp = boss_hp
                skill_player_hp = new_player_hp
                
                if skill_id == 0:  # 普通攻击
                    skill_boss_hp -= 5
                    skill_cooldown = 0
                elif skill_id == 1:  # 大招
                    skill_boss_hp -= 12
                    skill_cooldown = 2
                else:  # 治疗
                    skill_player_hp = min(100, skill_player_hp + 10)
                    skill_cooldown = 1
                
                skill_new_cooldowns = new_cooldowns.copy()
                skill_new_cooldowns[skill_id] = skill_cooldown
                
                # 更新技能序列
                new_sequence = skill_sequence + [skill_id]
                
                # 计算估计总回合数
                estimated_total_turns = turns_used + 1 + self.estimate_remaining_turns(
                    skill_boss_hp, skill_new_player_mp, skill_new_cooldowns
                )
                
                # 剪枝: 如果预估总回合数已经超过已知最优解，则跳过
                if estimated_total_turns >= self.best_score and self.best_score < float('inf'):
                    continue
                
                # 添加到队列
                heapq.heappush(priority_queue, (
                    estimated_total_turns,
                    turns_used + 1,
                    skill_boss_hp,
                    skill_player_hp,
                    skill_new_player_mp,
                    skill_new_cooldowns,
                    new_sequence
                ))
        
        return self.best_sequence

class BossBattle:
    def __init__(self, resources_from_maze=0):
        """
        初始化BOSS战斗
        resources_from_maze: 从迷宫带来的额外资源点数
        """
        self.boss = Boss()
        self.player = Player(initial_health=100 + resources_from_maze//2, 
                           initial_mp=30 + resources_from_maze//5)
        self.visualization = BossVisualization(self.boss, self.player)
        self.turn = 1
        self.max_turns = 20
        self.battle_end = False
        self.victory = False
        self.score = 0
        
        # 计算最优技能序列
        self.solver = BranchAndBoundSolver(
            self.boss.health, self.player.health, self.player.mp
        )
        self.optimal_sequence = self.solver.solve()
        self.current_optimal_step = 0

    def start_battle(self):
        """开始BOSS战斗"""
        self.visualization.add_battle_log("BOSS战斗开始!")
        
        if self.optimal_sequence:
            self.visualization.current_recommendation = self.optimal_sequence[0]
            self.visualization.recommended_sequence = self.optimal_sequence
            recommendations = [self.visualization.skills[skill_id]["name"] for skill_id in self.optimal_sequence]
            print(f"分支限界法推荐技能序列: {', '.join(recommendations)}")
        
        # 主循环
        while not self.battle_end:
            # 更新显示
            self.visualization.update_display()
            
            # 处理输入
            skill_choice = self.visualization.handle_events()
            
            if skill_choice is not None:
                self.execute_turn(skill_choice)
        
        # 战斗结束，显示结果
        self.visualization.update_display()
        time.sleep(1)
        
        # 显示最终结果
        if self.victory:
            self.score = self.calculate_score()
            result_text = f"恭喜击败BOSS! 用时{self.turn}回合，得分: {self.score}"
        else:
            result_text = "挑战失败!"
        
        self.visualization.add_battle_log(result_text)
        self.visualization.update_display()
        
        # 等待玩家点击关闭
        waiting = True
        while waiting:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    pygame.quit()
                    waiting = False
                elif event.type == pygame.KEYDOWN or event.type == pygame.MOUSEBUTTONDOWN:
                    waiting = False
        
        return self.score
    
    def execute_turn(self, skill_choice):
        """执行一个回合的战斗"""
        # 玩家回合
        result = self.player.use_skill(skill_choice, self.boss)
        self.visualization.add_battle_log(f"回合{self.turn}: {result}")
        
        # 检查BOSS是否被击败
        if self.boss.health <= 0:
            self.battle_end = True
            self.victory = True
            self.visualization.add_battle_log("你击败了BOSS!")
            return
        
        # BOSS回合
        damage = self.boss.attack_player(self.player)
        self.visualization.add_battle_log(f"BOSS对你造成{damage}点伤害!")
        
        # 检查玩家是否死亡
        if self.player.health <= 0:
            self.battle_end = True
            self.visualization.add_battle_log("你被BOSS击败了!")
            return
        
        # 回合结束，更新状态
        self.turn += 1
        self.player.regenerate_mp()
        self.player.update_cooldowns()
        
        # 检查是否超过最大回合数
        if self.turn > self.max_turns:
            self.battle_end = True
            self.visualization.add_battle_log("回合数已达上限，挑战失败!")
            return
        
        # 更新推荐技能
        if self.optimal_sequence and self.current_optimal_step < len(self.optimal_sequence) - 1:
            self.current_optimal_step += 1
            self.visualization.current_recommendation = self.optimal_sequence[self.current_optimal_step]
    
    def calculate_score(self):
        """计算最终得分"""
        return 1000 - 50 * self.turn + 2 * self.player.health + self.player.mp

    def attack_boss(self, damage):
        """
        攻击BOSS (保留旧有接口以兼容)
        """
        if damage > self.boss.defense:
            actual_damage = damage - self.boss.defense
            self.boss.health -= actual_damage
            self.visualization.add_battle_log(f"你对BOSS造成{actual_damage}点伤害!")
            if self.boss.health <= 0:
                self.visualization.add_battle_log("你击败了BOSS!")
                return True
        else:
            self.visualization.add_battle_log("你的攻击伤害太低!")
        
        return False
